from datetime import datetime
from common.log import logger
import time
import aiohttp
import asyncio
import traceback
from db.es.es_service import ElasticsearchStorage
from agent.llm_agent import LLMAgent
import settings

default_template = """
将以下内容归纳总结成一句话，字数限制在20个以内
===
{document}
===
        """


class ConversationToSummarize(LLMAgent):
    def __init__(self, document):
        self.prompt = default_template
        self.prompt = self.prompt.format(document=document)
        super().__init__(self.prompt)


class ConversationStoreService:
    def __init__(self, ):
        self.es = ElasticsearchStorage()

    def store_conversation_index(self, index_name, session_id: str, message_id: str, msg_type: str, text_info: str,
                                 attitude: str = 'no', role: str = 'assistant'):
        try:
            if session_id and session_id[:4] == "LXKJ":
                return
            conversation = {
                "session_id": session_id,
                "role": role,
                "message_id": message_id,
                "attitude": attitude,
                "msg_type": msg_type,
                "text_info": text_info,
                "message_time": datetime.now()
            }
            self.es.index_document(index_name, conversation)
        except Exception as ee:
            logger.error(traceback.format_exc())

    async def async_store_conversation_index(self, index_name, session_id: str, message_id: str, msg_type: str,
                                             text_info: str,
                                             attitude: str = 'no', role: str = 'assistant'):
        try:
            if session_id and session_id[:4] == "LXKJ":
                return
            conversation = {
                "session_id": session_id,
                "role": role,
                "message_id": message_id,
                "attitude": attitude,
                "msg_type": msg_type,
                "text_info": text_info,
                "message_time": datetime.now()
            }
            async with aiohttp.ClientSession():
                await self.es.index_document_async(index_name, conversation)
        except Exception as ee:
            logger.error(traceback.format_exc())

    def store_conversation_user(self, session_id, user_name, conversationHistory):
        # 用户<->会话关系存储
        session_title = None
        logger.info('save to es——{},{}'.format(user_name, session_id))
        if not conversationHistory:
            session_title = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        else:
            self.summarize_bot = ConversationToSummarize(f'''{conversationHistory}''')
            session_title = self.summarize_bot.chat_with_azure()
            logger.info('user：{},session_title:{}'.format(user_name, session_title))

        try:
            query = {
                "query": {
                    "bool": {
                        "must": [
                            {"term":
                                 {"session_id": session_id}
                             },
                            {"term":
                                 {"user_name": user_name}
                             }
                        ]
                    }
                }
            }
            results = self.es.search_documents(settings.USER_SESSION_INDEX, query)
            if not results:
                logger.info(
                    'create new user-sessionId-session_title：{}-{}'.format(user_name, session_id, session_title))
                # 创建用户会话
                conversation = {
                    "user_name": user_name,
                    "session_id": session_id,
                    "session_title": session_title,
                    "message_time": datetime.now()
                }
                self.es.index_document(settings.USER_SESSION_INDEX, conversation)
            elif conversationHistory and session_title:
                logger.info('update title')
                # 更新会话标题
                updated_fields = {
                    'session_title': session_title
                }
                self.es.update_document(settings.USER_SESSION_INDEX, results[0]['_id'], updated_fields)
        except Exception as ee:
            logger.error(traceback.format_exc())
